import torch
from transformers import BertTokenizer, BertModel
import numpy as np
import torch.nn.functional as F


def fine_grade_tokenize(raw_text, tokenizer):
    """
    序列标注任务 BERT 分词器可能会导致标注偏移，
    用 char-level 来 tokenize
    """
    tokens = []

    for _ch in raw_text:
        if _ch in [' ', '\t', '\n']:
            tokens.append('[BLANK]')
        else:
            if not len(tokenizer.tokenize(_ch)):
                tokens.append('[INV]')
            else:
                tokens.append(_ch)

    return tokens


def calc_similarity(s1, s2):
    # 计算余弦相似度，并返回结果
    cosine_sims = F.cosine_similarity(s1.mean(dim=1), s2.mean(dim=1), dim=1)
    # sim = np.dot(s1.detach().numpy(), s2.T.detach().numpy()) / (np.linalg.norm(s1.detach().numpy()) * np.linalg.norm(s2.detach().numpy()))
    # print(sim.sum())
    return cosine_sims.mean().detach().numpy()


def bert_encoding(sentence):
    tokens = fine_grade_tokenize(sentence, bert_tokenizer)
    encode_dict = bert_tokenizer.encode_plus(text=tokens, is_pretokenized=True, return_token_type_ids=True,
                                             return_attention_mask=True, max_length=128,
                                             pad_to_max_length=True)  # max_length=1000, pad_to_max_length=True
    token_ids = torch.tensor(encode_dict['input_ids'])
    attention_masks = torch.tensor(encode_dict['attention_mask'])
    token_type_ids = torch.tensor(encode_dict['token_type_ids'])
    zip_result = list(zip(tokens, token_ids[1:-1].numpy()))

    # 修改维度，用于接下来的Bert编码
    token_ids = token_ids.unsqueeze(0)
    attention_masks = attention_masks.unsqueeze(0)
    token_type_ids = token_type_ids.unsqueeze(0)
    bert_outputs = bert_module(input_ids=token_ids, attention_mask=attention_masks, token_type_ids=token_type_ids)
    seq_out = bert_outputs[0]
    return seq_out


def bert_encoding(sentence):
    word_embeddings_list = []
    for text in sentence:
        encoded_input = bert_tokenizer(text, return_tensors='pt')
        outputs = bert_module(**encoded_input)
        # print(outputs.last_hidden_state.shape)
        word_embeddings = outputs.last_hidden_state[0, :-1, :]
        word_embeddings_list.append(word_embeddings)

    tokens = fine_grade_tokenize(sentence, bert_tokenizer)
    encode_dict = bert_tokenizer.encode_plus(text=tokens, is_pretokenized=True, return_token_type_ids=True,
                                             return_attention_mask=True, max_length=128,
                                             pad_to_max_length=True)  # max_length=1000, pad_to_max_length=True
    token_ids = torch.tensor(encode_dict['input_ids'])
    attention_masks = torch.tensor(encode_dict['attention_mask'])
    token_type_ids = torch.tensor(encode_dict['token_type_ids'])
    zip_result = list(zip(tokens, token_ids[1:-1].numpy()))

    # 修改维度，用于接下来的Bert编码
    token_ids = token_ids.unsqueeze(0)
    attention_masks = attention_masks.unsqueeze(0)
    token_type_ids = token_type_ids.unsqueeze(0)
    bert_outputs = bert_module(input_ids=token_ids, attention_mask=attention_masks, token_type_ids=token_type_ids)
    seq_out = bert_outputs[0]
    return seq_out


# bert_dir =  "/home/Dyf/code/models/pretrain_models/bert"
bert_dir = "/home/Dyf/code/models/pretrain_models/bertcn/bert-base-chinese"
bert_tokenizer = BertTokenizer.from_pretrained(bert_dir)
bert_module = BertModel.from_pretrained(bert_dir)
word2id_bert = bert_tokenizer.get_vocab()
id2word_bert = {word2id_bert[key]: key for key in word2id_bert}

# s1 = "早泄症状早期症状：精神抑郁、头晕、神疲乏力，在接触或刚接触女方外阴，或时间短暂，尚未进入高潮即射精，还伴有长期精神不振，夜寐不安，精薄清冷等情况。晚期症状：射精过快，插入后，在女士尚未达到性高潮，而男性的时间短于2分钟，提早射精而出现的不和谐性障碍。相关症状：器质性早泄男子性功能障碍一、早泄症状早泄的临床表现主要是射精过快。典型症状是指插入后，在女士尚未达到性高潮，而男性的时间短于2分钟，提早射精而出现的不和谐障碍，一般男性30%均有此情况。一般以为，早泄是指男子在勃起之后，未进入之前，或正当纳入、以及刚刚进入而尚未抽动时便已射精，也自然随之疲软并进入不应期的现象。临床上对勃起未进入即射精，诊断为早泄。男性初期早泄症状主要是精神抑郁、头晕、神疲乏力，在接触或刚接触女方外阴，或时间短暂，尚未进入高潮即射精，还伴有长期精神不振，夜寐不安，精薄清冷等情况。二、早泄判断标准早泄的判断标准是什么？早泄可以说是先在社会中男性常见的一种现象，不规律的生活习惯等都可能导致得了早泄。偶尔的一次射精过快很多人就觉得得了早泄，其实不然有时候早泄也是会被男性自身错误判断给误导，因而导致心情闷闷不乐，对生活、工作都套上了阴影开始自惭形愧。早泄的判断标准是什么？专家表示判断早泄的标准主要看以下几方面：1、以时间为标准从阴茎插入阴道至射精的时间，一般认为短于2分钟即为早泄，但严格者认为应短于30秒才能算早泄。2、以抽动次数为标准阴茎插入阴道中抽动次数少于10次至30次为早泄。3、以性伴侣的反应为标准在性活动中，如果有半数以上的性生活机会中，不能使女方达到性高潮亦可称为早泄。4、以控制射精反射的能力为标准在性交时射精过快，可于当晚再次性交，如射精时间相对较长，可不必担忧。这种重复性交亦可作为对待早泄的偶一为之的办法，但不宜常规使用。另有一种情况就是由于男子性激动，尿道口分泌出一些液体，其实是一种尿道球腺等腺体分泌的粘液，不可误认为是早泄的精液。过度度疲劳时性交偶有一次过早射精，新婚或久别重逢第一次性交，男方过于激动，往往射精较快。偶尔一次并不是病态，经常如此，才能说是早泄。所以当男性偶尔出现早泄症状时千万不要着急，要调整好心态相信自己。但是经常如此，那么就是得了早泄了。分居两地的夫妇，相逢第一晚男方偶然发生了早泄，思想负担太重，在大脑皮层里留下了不良刺激痕迹，以后每次久别重逢第一夜过性生活时都会发生早泄，如果二、三晚连续出现，这种特殊现象，就属于病态早泄。一些人入睡前性交容易发生过快射精，而下半夜睡醒之后性交，则射精时间较长，这种情况不应视为早泄，可以因事之宜而为之。国际性医学会（InteRNAtionalSocietyforSexualMedicineISSM）从循证医学的角度上指出早泄的定义应包括以下三点：①射精总是或者几乎总是发生在阴茎插入阴道1分钟以内；②不能全部或几乎全部进入阴道后延迟射精；③消极的个人精神心理因素，比如苦恼、忧虑、挫折感和/或逃避性活动等。三、分类1.原发性早泄原发性早泄少见，难以诊断。特点是：(1)第一次性交出现；(2)对性伴侣，没有选择性；(3)每次性交都发生过早射精。2.继发性早泄继发性早泄是后天获得的早泄，有明确的生理或者心理病因。特点是：(1)过早射精发生在一个明确的时间；(2)发生过早射精前射精时间正常；(3)可能是逐渐出现或者突然出现：(4)可能继发于泌尿外科疾病、甲状腺疾病[9]或者心理疾病等。3.境遇性早泄国内也有学者将此类早泄称为自然变异性早泄。此类患者的射精时间有长有短，过早射精时而出现。这种早泄不一定是都病理过程，具体特点是：(1)过早射精不是持续发生，发生时间没有规律。(2)在将要射精时，控制射精的能力降低，但有时正常，这点不是诊断的必要条件。4.早泄样射精功能障碍此类患者射精潜伏时间往往在正常范围，患者主观上认为自己早泄，此类早泄不能算是真正的病理过程，通常隐藏着心理障碍或者与性伴侣的关系问题。此类早泄的特点是：(1)主观认为持续或者非持续射精过快；(2)患者自己想象中的过早射精或者不能控制射精焦虑；(3)实际插入阴道射精潜伏时间正常甚至很长；(4)在将要射精时，控制射精的能力变低；(5)用其他精神障碍不能解释患者的焦虑。还是勃起功能障碍，当然，也有不少病人既有早泄又有勃起功能障碍。"
# s2 = "阳痿症状典型症状：男性在性交开始阶段勃起缓慢或者难以勃起；有时虽然勃起了，但勃起的硬度不够而造成插入困难或者根本不能插入阴道。相关症状：心理性性功能障碍男子性功能障碍雄激素过少阳痿症状主要表现为：男性在性交开始阶段勃起缓慢或者难以勃起；有时虽然勃起了，但勃起的硬度不够而造成插入困难或者根本不能插入阴道。男性朋友如果身体有下面这些表现时，就要注意是否患上ED了。性欲障碍：包括无性欲、性欲低下、性厌恶、性欲亢进等；阴茎勃起障碍：包括阳痿、阴茎勃起不坚、阴茎异常勃起等；性交障碍：包括性交昏厥、性交失语、性交癔病、性交恐惧症等；射精障碍：包括早泄、遗精、不射精、射精延迟、逆行射精、射精疼痛、血精等。上述这几个方面可以单独出现，也可多个同时出现，其中最多见的是阴茎勃起障碍和射精障碍。需要强调的是，正常男性的性功能也存在着生理性的波动。当性功能在精神、情绪不稳定，疲劳，健康状况不佳或女方对性生活冷淡或持反对态度等因素刺激时，均可出现一时性的“阳痿”，这种偶然现象不能视为病态。只有在排除上述诸因素的影响，在正常性刺激下，反复多次出现性交失败，方能认为是阳痿。如果一旦患了阳痿或怀疑自己患了本病，就应该及时到正规医院找专科医生诊治，且勿讳疾忌医，贻误病机。"

s1 = """
工业4.0技术改变了各研究领域的传统方法,利用数字化实现可持续性,并引入创新基础设施。

在当今时代,每个组织都需要一个独特的营销策略,以满足客户和市场需求的产品和服务的形式。

客户满意度、客户保留、客户行为、客户画像和奖励系统是有效实施组织营销策略以实现更好的生产力的关键参数。

有有限的研究集中于讨论营销中用于改造数字和智能生态系统的所有工业4.0使能技术。

在分析的基础上,本研究确定了工业4.0使能技术在市场策略中的应用,例如目标客户的客户满意度战略信息；开发用于接收产品和服务实时反馈的数字基础设施；预测客户行为,制定个性化信息或服务；利用商业分析来加强产品或服务的质量；根据消费者和市场需求,开发有效的模拟来监测、测试和计划产品改进。

最后,建议了一个框架,并讨论了在保持可持续性的同时,对未来采用的重要建议。 精简提炼观点
"""
s2 = """
1. 工业4.0技术正在改变研究的领域，促进数字化的实现，以及推动创新基础设施的引入。
2. 在当前的环境下，每个组织都需要特定的市场策略以提供符合客户和市场需要的产品与服务。
3. 为了提高组织的生产力，需关注包括客户满意度，客户维持，客户行为，客户画像和奖励系统等重要因素。
4. 对于如何使用工业4.0技术改革市场营销策略方面存在研究缺口。
5. 本研究分析了工业4.0技术如何应用于市场策略，如通过收集实时反馈的信息改善产品和服务，或者使用商业分析提升产品的品质等。
6. 推荐在未来考虑可持续性的同时，使用一个框架去适应工业4.0的技术发展。
"""


def split_sentences(s1):
    split_sentences = []
    n_split = s1.split("\n")
    for i in n_split:
        juhao_split = i.split()
        split_sentences.extend(juhao_split)
    split_sentences = [i for i in split_sentences if i]
    return split_sentences


s1_split = split_sentences(s1)
s2_split = split_sentences(s2)

for i in s1_split:
    print("############", i)
    for j in s2_split:
        print("$$$$$$$$$$$$", j)
        s1_out = bert_encoding(i)
        s2_out = bert_encoding(j)
        similarity = calc_similarity(s1_out, s2_out)
        print(f"相似度：{similarity:.4f}")
